import json
import numpy as np
import torch
import torch.utils.data as Data

# 创建训练数据集, 从原始训练文件中将中文字符进行数字编码, 并将标签页进行数字编码
def create_train_data(train_data_file, result_file, json_file, tag2id, max_length=20):
    # 导入json格式的中文字符到id的映射表
    char2id = json.load(open(json_file, mode='r', encoding='utf-8'))

    char_data, tag_data = [], []

    # 打开原始训练文件
    with open(train_data_file, mode='r', encoding='utf-8') as f:
        # 初始化一条语句数字化编码后的列表
        char_ids = [0] * max_length
        tag_ids = [0] * max_length
        idx = 0
        for line in f.readlines():
            line = line.strip('\n').strip()
            # 如果不是空行, 并且当前语句长度没有超过max_length, 则进行字符到id的映射
            if len(line) > 0 and line and idx < max_length:
                ch, tag = line.split('\t')
                # 如果当前字符存在于映射表中, 则直接映射为对应的id值
                if char2id.get(ch):
                    char_ids[idx] = char2id[ch]
                # 否则直接用"UNK"的id值来代替这个未知字符
                else:
                    char_ids[idx] = char2id['UNK']
                # 将标签也进行对应的转换
                tag_ids[idx] = tag2id[tag]
                idx += 1
            # 如果是空行, 或者当前语句长度超过max_length
            else:
                # 如果当前语句长度超过max_length, 直接将[0: max_langth]的部分作为结果
                if idx <= max_length:
                    char_data.append(char_ids)
                    tag_data.append(tag_ids)
                # 遇到空行, 说明当前句子已经结束, 初始化清零, 为下一个句子的映射做准备
                char_ids = [0] * max_length
                tag_ids = [0] * max_length
                idx = 0

    # 将数字化编码后的数据封装成numpy的数组类型, 数字编码采用np.int32
    x_data = np.array(char_data, dtype=np.int32)
    y_data = np.array(tag_data, dtype=np.int32)

    np.savez(result_file, x_data=x_data, y_data=y_data)  # 直接利用np.savez()将数据存储为.npz类型的文件
    print("create_train_data Finished!".center(100, "-"))


# 生成批量训练数据
def load_dataset(data_file, batch_size):
    # 将第二步生成的train.npz文件导入内存
    data = np.load(data_file)

    # 分别取出特征值和标签
    x_data = data['x_data']
    y_data = data['y_data']

    # 将数据封装成tensor张量
    x = torch.tensor(x_data, dtype=torch.long)
    y = torch.tensor(y_data, dtype=torch.long)

    # 将数据封装成Tensor数据集
    dataset = Data.TensorDataset(x, y)

    total_length = len(dataset)

    # 采用80%的数据作为训练集, 20%的数据作为测试集
    train_length = int(total_length * 0.8)
    validation_length = total_length - train_length

    # 利用Data.random_split()直接切分集合, 按照80%, 20%的比例划分
    train_dataset, validation_dataset = Data.random_split(dataset=dataset,
                                        lengths=[train_length, validation_length])

    # 将训练集进行DataLoader封装
    # 参数说明如下:
    # dataset:     训练数据集
    # batch_size:  代表批次大小, 若数据集总样本数量无法被batch_size整除, 则最后一批数据为余数
    #              若设置drop_last为True， 则自动抹去最后不能被整除的剩余批次
    # shuffle:     是否每个批次为随机抽取, 若为True, 则每次迭代时数据为随机抽取
    # num_workers: 设定有多少子进程用来做数据加载, 默认为0, 即数据将被加载到主进程中
    # drop_last:   是否去除不能被整除后的最后批次, 若为True, 则不生成最后不能被整除剩余的数据内容
    #              例如: dataset长度为1028, batch_size为8,
    #              若drop_last=True, 则最后剩余的4(1028/8=128余4)条数据将被抛弃不用
    train_loader = Data.DataLoader(dataset=train_dataset, batch_size=batch_size,
                                   shuffle=True, num_workers=4, drop_last=True)

    validation_loader = Data.DataLoader(dataset=validation_dataset, batch_size=batch_size,
                                        shuffle=True, num_workers=4, drop_last=True)

    # 将两个数据生成器封装为一个字典类型
    data_loaders = {'train': train_loader, 'validation': validation_loader}

    # 将两个数据集的长度也封装为一个字典类型
    data_size = {'train': train_length, 'validation': validation_length}

    return data_loaders, data_size



# 参数1:字符码表文件路
json_file = 'char_to_id.json'

# 参数2:标签码表对照字典
tag2id = {"O": 0, "B-dis": 1, "I-dis": 2, "B-sym": 3, "I-sym": 4, "<START>": 5, "<STOP>": 6}

# 参数3:训练数据文件路径
train_data_file = 'train.txt'

# 参数4:创建的npz文件保路径(训练数据)
result_file = 'train.npz'

# 批次大小
batch_size = 8


if __name__ == '__main__':
    create_train_data(train_data_file, result_file, json_file, tag2id)
    data_loader, data_size = load_dataset(result_file, batch_size)
    print('data_loader:', data_loader, '\ndata_size:', data_size)